#include <unistd.h>
#include <fcntl.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/types.h>

#include "TCPNetwork.hpp"
#include "../constants.hpp"
//#include "../../membership/Principal.hpp"


using namespace BFT::network;

TCPNetwork::~TCPNetwork()
{
  for(std::vector<listenerInfo>::iterator iter = listenerSockets.begin();
      iter != listenerSockets.end(); iter++) {
    close(iter->listenerSocket);
  }
}

TCPNetwork::TCPNetwork(BFT::util::Role _role, boost::shared_ptr<BFT::membership::Membership> _membership)
{
  myRole = _role;
  membership = _membership;
  listening = false;
  myid = membership.get()->getMyId();
  justaccepted=false;
	
  max_fd = -1;
  FD_ZERO(&rdfs);
	
  //myInterfaces is a shared pointer to a vector of shared pointers to principals
  boost::shared_ptr< std::vector< boost::shared_ptr<BFT::membership::Principal> > > myInterfaces 
    = membership.get()->getMyInterfaces(myRole);
	 
  int interfacesNumber = myInterfaces.get()->size();
	
  //fprintf(stderr,"I have %d interfaces\n",interfacesNumber);
	
  //reserving anough space in the vector for all the sockets
  //listenerSockets.reserve(interfacesNumber);
	
  //creating a server socket for each interface in myinterfaces
  for (int index=0; index < interfacesNumber ; index++)
    {
      //initialize socket
      listenerInfo tmpinfo;
      tmpinfo.listenerSocket = socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);
      if (tmpinfo.listenerSocket < 0)
	{
	  fprintf(stderr,"unable to initialize socket. index %d\n",index);
	  exit(24);
	}
	 	
      //create the address
      tmpinfo.socketAddress.sin_family = AF_INET;
      tmpinfo.socketAddress.sin_port = htons(myInterfaces.get()->at(index).get()->getPort());
      //	 	tmpinfo.socketAddress.sin_addr.s_addr = myInterfaces.get()->at(index).get()->getAddr();
      tmpinfo.socketAddress.sin_addr.s_addr = INADDR_ANY;
	 	
      //binding the socket
      if ( bind( tmpinfo.listenerSocket , (struct sockaddr*) &tmpinfo.socketAddress , sizeof(tmpinfo.socketAddress) ) < 0 )
	{
	  fprintf(stderr,"unable to bind socket. index %d, port: %d, errno:%d\n"
		  ,index,myInterfaces.get()->at(index).get()->getPort(), errno );
	  exit(25);
	}
	 	
      //making the sokcet asynchronous
      int currentflags = fcntl(tmpinfo.listenerSocket , F_GETFL);
      if (currentflags < 0)
	{
	  fprintf(stderr,"unable to get flags from socket. index%d\n",index);
	  exit(26);
	}
	 	
      if (  fcntl(tmpinfo.listenerSocket, F_SETFL, currentflags | O_NONBLOCK) < 0 )
	{
	  fprintf(stderr,"unable to set flags from socket. index%d\n",index);
	  exit(27);
	}
	 	
      // setting SO_REUSEADDR option
      int yes=1;
      setsockopt(tmpinfo.listenerSocket, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int));	 	
	 	
      //setting the socket for listening
      if ( listen( tmpinfo.listenerSocket, BACKLOG_SIZE ) < 0 )
	{
	  fprintf(stderr,"unable to set socket for listening, index%d\n",index);
	  exit(28);
	}		 	
	 	
      //register this listener for select
      FD_SET(tmpinfo.listenerSocket, &rdfs);
      max_fd = std::max(max_fd, tmpinfo.listenerSocket);
	 		
      //add <listener, index> to indexTable
      listenerIndex[tmpinfo.listenerSocket] = index;	
      listenerSockets.push_back(tmpinfo);
    }
	
  //now:
  //1- all my interfaces are associated with a listener
  //2- all the listeners are registered in the fd_set rdfs
  //3- all the listeners can be used as key to retreive the index of the corresponding principal
	 
}


void TCPNetwork::acceptConnection(int _listener)
{
  //what happened before:
  //select() told us that there is something to read on
  // a file descriptor AND the file descriptor belong to the
  // listner_set -> accepting a new connection
	
  socketInfo tempinfo;

  tempinfo.addrlen = sizeof(tempinfo.remote_address);
  tempinfo.socketFd = accept( _listener, &(tempinfo.remote_address), &(tempinfo.addrlen)); 
	
  if (tempinfo.socketFd < 0)
    {
      fprintf(stderr, "unable to accept\n");
      exit(29);
    }
		
  //complete tempinfo
  tempinfo.interfaceIndex = listenerIndex[_listener];
  tempinfo.message_length = -1; 
  tempinfo.bytes_size = 0;
	
  //put tempinfo inside the table of sockets
  socketTable.insert(std::pair<int, socketInfo> (tempinfo.socketFd, tempinfo));
	
  //	fprintf(stderr,"New connection accepted, %d connections concurrently open\n", (int)socketTable.size());
	
  //make the socket non_blocking
  int currentflags = fcntl(tempinfo.socketFd , F_GETFL);
  if (currentflags < 0)
    {
      fprintf(stderr,"unable to get flags from socket. index%d\n",listenerIndex[_listener]);
      exit(30);
    }
 	
  if (  fcntl(tempinfo.socketFd, F_SETFL, currentflags | O_NONBLOCK) < 0 )
    {
      fprintf(stderr,"unable to set flags from socket. index%d\n",listenerIndex[_listener]);
      exit(31);
    }
	
  justaccepted=true;	
}

//
boost::tuple<int,boost::shared_ptr<char>,int>TCPNetwork::readMessage(int _socket)
{
  //what happened before:
  // select() told us that there is something to read on
  // the file descriptor _socket AND _socket does *not*
  // belong to the listener_set -> _socket is the descriptor
  // associated to an already accepted connection -> I expect
  // to receive messages
	
  //2 possible cases:
  //1 - no partial messages saved for this connection, preparing to receive a new message
  //2 - already have a partial message
	
  //retrieve the socketInfo related to  _socket
  socketInfo tempInfo = socketTable[_socket];
	
  //boost::tuple<int,boost::shared_ptr<char>,int> empty;
	
  char marker [MARKER_SIZE];
	
  if (tempInfo.message_length == -1)
    {
      //no partial messages cached, preparing to receive the beginning of a message, preceded
      // (and followed) by  the 4 bytes marker conteining the message size
      bzero(marker,MARKER_SIZE);
      //peeking the header, to see if we can read a full one
      int numPeekedBytes = recv(_socket, (void*)marker, MARKER_SIZE, MSG_PEEK);
      if (numPeekedBytes != MARKER_SIZE)
	{
	  if ( (numPeekedBytes < 1) && (errno != EAGAIN) )
	    {
	      fprintf(stderr,"unexpected error, removing socket\n");
	      socketToErase.push_back(_socket);
	    } 
	  else
	    {	
	      if(numPeekedBytes == 0)
		{
		  fprintf(stderr,"connection closed by remote endpoint, removing socket\n");
		  socketToErase.push_back(_socket);
		}
	    }
								
	  //not even able to read a full marker, returning an empty tuple
	  boost::shared_ptr<char> i;
	  return boost::tuples::make_tuple(0, i ,0);
	}
		
      //so, it is possible to read a complete marker, than read it!
      int numReadBytes = recv(_socket, (void *)marker, MARKER_SIZE, 0);
		
      if (numReadBytes != MARKER_SIZE)
	{	
	  //peeked 4 bytes, have to be able to read them
	  fprintf(stderr,"weird error occurred... read only %d bytes after a peek of %d",numReadBytes, MARKER_SIZE);
	  exit(40);
	}
		
      //initialize tempinfo data structure
      tempInfo.message_length = ntohl( *(uint32_t*)marker ); //setting expected message length
      tempInfo.bytes.reset( (char*)malloc(tempInfo.message_length) , free ); //initializing shared_ptr
      tempInfo.bytes_size = 0; //I currently have 0 bytes
		
      //updating socket information
      socketTable[_socket]= tempInfo;
    }
	
  //we have a cached message. either:
  //1- the message is not complete
  //2- the message is complete, and we have not cleaned the tail marker yet
	
  if (tempInfo.bytes_size < tempInfo.message_length)
    {
		
      // message not complete. read from the socket to fill the bytes

      int numReadBytes = recv(_socket, (void*)(tempInfo.bytes.get()+tempInfo.bytes_size)
			      , tempInfo.message_length - tempInfo.bytes_size,0);

      if ( numReadBytes < 1 )
	{
	  if (errno != EAGAIN)
	    {
	      fprintf(stderr,"unexpected error, removing socket\n");
	      socketToErase.push_back(_socket);
	    }
	  boost::shared_ptr<char> i;
	  return boost::tuples::make_tuple(0, i ,0);
	} 
		
	
      if(numReadBytes == 0)
	{
	  fprintf(stderr,"connection closed by remote endpoint, removing socket\n");
	  socketToErase.push_back(_socket);
	  boost::shared_ptr<char> i;
	  return boost::tuples::make_tuple(0, i ,0);
	}

      //no errors occurred
		
      if ( numReadBytes < tempInfo.message_length - tempInfo.bytes_size )
	{
	  //read something, but the message is not complete yet
	  //updating socket information and returning empty tuple
	  tempInfo.bytes_size += numReadBytes;
	  socketTable[_socket]=tempInfo;
	  boost::shared_ptr<char> i;
	  return boost::tuples::make_tuple(0, i ,0);
			
	}
		
      if ( numReadBytes == tempInfo.message_length - tempInfo.bytes_size )
	{
	  //read the whole message content
	  //(though the tail marker is still to be received)
	  tempInfo.bytes_size = tempInfo.message_length;
	  socketTable[_socket] = tempInfo;
	}
    }
	
  //so, if we made it through there, we have a complete message inside tempInfo.bytes. 
  //we only need to read the tail marker before cleaning the socket and adding a 
  //complete message to the vector
	
  //just a safety check
  if (tempInfo.bytes_size != tempInfo.message_length)
    {
      fprintf(stderr,"if that happened, I've made something horribly wrong...\n");
      exit(31);
    }
	
  bzero(marker,MARKER_SIZE);
  //try to read a full marker
  int numPeekedBytes = recv(_socket, (void*)marker, MARKER_SIZE, MSG_PEEK);
  if (numPeekedBytes != MARKER_SIZE)
    {
      if ( (numPeekedBytes < 1) && (errno != EAGAIN) )
	{
	  fprintf(stderr,"unexpected error, removing socket\n");
	  socketToErase.push_back(_socket);
	} 
      else
	{	
	  if(numPeekedBytes == 0)
	    {
	      fprintf(stderr,"connection closed by remote endpoint, removing socket\n");
	      socketToErase.push_back(_socket);
	    }
	}
							
      //not even able to read a full marker, returning an empty tuple
      boost::shared_ptr<char> i;
      return boost::tuples::make_tuple(0, i ,0);
    }
	
  //it is possible to read a complete marker, than read it!
  int numReadBytes = recv(_socket, (void *)marker, MARKER_SIZE, 0);
  if (numReadBytes != MARKER_SIZE)
    {	
      fprintf(stderr,"weird error occurred... read only %d bytes after a peek of %d",numReadBytes, MARKER_SIZE);
      exit(41);
    }

  //check that it really is the tail marker
  if ( (uint32_t)tempInfo.message_length != ntohl( *(uint32_t*)marker ) )
    {
      fprintf(stderr,"either something wrong in message handling OR malformed message from client\n");
      fprintf(stderr,"for now, I am assuming I did something wrong.\n");
      exit(32);
      //TODO
      //if it is client's fault, I probably want to close the connection
    }
	
  //ok, read a complete message and cleaned up the tail marker,
  //building the tuple to be returned and cleaning up the socket state
	
  boost::tuples::tuple <int, boost::shared_ptr<char>, int> to_return (tempInfo.interfaceIndex , tempInfo.bytes , tempInfo.bytes_size);

  tempInfo.bytes.reset();
  tempInfo.bytes_size=0;
  tempInfo.message_length=-1;
	
  socketTable[_socket]=tempInfo;
  return to_return;
}

std::vector< boost::tuple< int , boost::shared_ptr<char> , int > > TCPNetwork::readMessages(int _socket, int _max)
{
  //just loop around readMessage() and stop when an empty message is returned
  // or when _max messages have already been read
	
  std::vector< boost::tuple< int , boost::shared_ptr<char> , int > > to_return;
	
  boost::tuple< int , boost::shared_ptr<char> , int > tempTuple;
	
  tempTuple = readMessage(_socket);
	
  while( ( tempTuple.get<2>() > 0) && (to_return.size() < (uint32_t)_max) )
    {
      to_return.push_back(tempTuple);
      tempTuple = readMessage(_socket);			
    }
	
  return to_return;	
}


std::vector< boost::tuple< int , boost::shared_ptr<char> , int > >TCPNetwork::selectMessages()
{
  int ready;
  std::vector< boost::tuple <int, boost::shared_ptr<char> , int> > msg_vector;
  boost::tuple <int, boost::shared_ptr<char> , int> message;
	
  //keep copy of the current rdset
  fd_set previous;
		
  std::map<int, int>::iterator listenerIterator;
  std::map<int, socketInfo>::iterator socketIterator;
	
  while(true)
    {
      timeout.tv_sec=TV_SEC;
      timeout.tv_usec=TV_USEC;
		

      memcpy((void*)&previous, (const void*)&rdfs, sizeof(rdfs));
      ready = select(max_fd+1, &rdfs, NULL, NULL, &timeout );

      if(ready == 0)
	{
	  memcpy((void*)&rdfs, (const void*)&previous, sizeof(rdfs) );
	  continue;
	}

      if (ready <0)
	{
	  fprintf(stdout,"error in select, errno:%d\n",errno);
	  exit(50);
	  //TODO
	  //should handle at least EINTR...
	}
      /*		
			if (ready == 0)
			fprintf(stderr, "nothing interesting happening there......\n");
      */
      if (ready > 0)
	{

	  //something ready to be read, either:
	  // 1- is a socket listening for incoming connection -> call accept
	  // 2- is a socket related to a connection already established -> try to receive a message
			
	  for( listenerIterator = listenerIndex.begin() //initialize iterator
		 ; (listenerIterator != listenerIndex.end()) && (ready > 0) //run through the map, but stop if ready becomes 0
		 ; listenerIterator++)
	    {
	      if (FD_ISSET(listenerIterator->first,&rdfs))
		{
		  //					fprintf(stderr,"accepting a new connection....\n");
		  //a socket in listening mode have been seleted, calling accept
		  acceptConnection(listenerIterator->first);
		  ready--;
		} 
	    }
			
	  for( socketIterator = socketTable.begin() //initialize iterator
		 ; (socketIterator != socketTable.end()) && (ready > 0) //run through the map, but stop if ready becomes 0
		 ; socketIterator++)
	    {
	      if (FD_ISSET(socketIterator->first,&rdfs))
		{
		  //at least a message to read 
		  message=readMessage(socketIterator->first);
		  //add message to the vector if not empty
		  if (message.get<2>() != 0)
		    {
		      //fprintf(stderr,"read a message of size %d\n",message.get<2>());
		      msg_vector.push_back(message);
		    }
		  ready--;
		} 
	    }
	}
		

      if(socketToErase.empty() && !justaccepted)
	{
	  // still listening on the same socket as before
	  //fprintf(stderr, "lazy memcopy\n");
	  memcpy((void*)&rdfs, (const void*)&previous, sizeof(rdfs) );
	}
      else
	{
	  //fprintf(stderr, "rebuilding rdfs\n");
	  //we closed some connections, rebuilding the readset
	  FD_ZERO(&rdfs);
	  max_fd=-1;
	  for( listenerIterator = listenerIndex.begin() //initialize iterator
		 ; listenerIterator != listenerIndex.end()
		 ; listenerIterator++)
	    {
	      FD_SET(listenerIterator->first, &rdfs);
	      max_fd = std::max(max_fd, listenerIterator->first);
	    }
			
	  //TODO: optimize that...	
	  while(!socketToErase.empty())
	    {
	      socketTable.erase( socketToErase.back() );
	      socketToErase.pop_back();
	      /*
		if(socketToErase.empty())
		fprintf(stderr,"deleted one or more sockets, %d sockets concurrently open\n",(int)socketTable.size());
	      */	
	    }


	  for( socketIterator = socketTable.begin() //initialize iterator
		 ; socketIterator != socketTable.end()
		 ; socketIterator++)
	    {	
	      FD_SET(socketIterator->first, &rdfs);
	      max_fd = std::max(max_fd, socketIterator->first);
	    }
			
	  justaccepted=false;
	}

      if(msg_vector.size() >0)
	{
	  return msg_vector;
	}
		
    }// end of the selection loop.
}

void TCPNetwork::start()
{
  {
    boost::unique_lock<boost::mutex> lock(condition_mutex);
    listening = true;
  }	
  listening_condition.notify_all();
}

void TCPNetwork::stop()
{
  fprintf(stderr,"not implemented yet\n");
  exit(65);
}

void TCPNetwork::waitForListening()
{
  boost::unique_lock<boost::mutex> lock(condition_mutex);
  while( !listening)
    {
      listening_condition.wait(lock);
    }
}

/////////////////////////////////////////////////////





